# captcha_ctc_train.py

import os
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image

# 字符集和映射
CHARS = list(
    "0123456789一二三四五六七八九壹贰叁肆伍陆柒捌玖加减ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz家见等问")
MAX_LABEL_LEN = 7
BLANK_LABEL = 0  # CTC blank

char2idx = {c: i + 1 for i, c in enumerate(CHARS)}  # 0 reserved for blank
idx2char = {i + 1: c for i, c in enumerate(CHARS)}
idx2char[BLANK_LABEL] = ''


# -------------------------
# 数据集定义
# -------------------------
class CaptchaDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        self.image_dir = image_dir
        self.filenames = [f for f in os.listdir(image_dir) if f.endswith('.png')]
        self.transform = transform or transforms.Compose([
            transforms.Grayscale(),
            transforms.Resize((40, 100)),
            transforms.ToTensor(),
        ])

    def __len__(self):
        return len(self.filenames)

    def encode_label(self, label):
        return torch.tensor([char2idx.get(c, 0) for c in label], dtype=torch.long)

    def __getitem__(self, idx):
        fname = self.filenames[idx]
        label_text = fname.split('_')[0]
        image_path = os.path.join(self.image_dir, fname)
        image = Image.open(image_path).convert('L')

        image = self.transform(image)
        label = self.encode_label(label_text)
        return image, label, len(label)


# -------------------------
# 模型定义（CNN + BiLSTM + CTC）
# -------------------------
class CRNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),  # 20x80
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),  # 10x40
        )
        self.rnn = nn.LSTM(input_size=64 * 10, hidden_size=128, num_layers=2, bidirectional=True, batch_first=True)
        self.fc = nn.Linear(128 * 2, num_classes)

    def forward(self, x):
        b, c, h, w = x.size()
        x = self.cnn(x)
        x = x.permute(0, 3, 1, 2).contiguous()  # [B, W, C, H]
        x = x.view(b, x.size(1), -1)  # [B, W, C*H]
        x, _ = self.rnn(x)
        x = self.fc(x)
        return x.log_softmax(2)  # [B, T, C]


# -------------------------
# 训练流程
# -------------------------
def collate_fn(batch):
    images, labels, label_lens = zip(*batch)
    images = torch.stack(images)
    targets = torch.cat(labels)
    label_lens = torch.tensor(label_lens, dtype=torch.long)
    input_lens = torch.full(size=(images.size(0),), fill_value=25, dtype=torch.long)  # 修改为固定25
    # 或根据模型动态获取
    return images, targets, input_lens, label_lens


def train():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train_dataset = CaptchaDataset('/temp/captcha/train/')
    val_dataset = CaptchaDataset('/temp/captcha/test/')
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)

    model = CRNN(num_classes=len(CHARS) + 1).to(device)
    criterion = nn.CTCLoss(blank=BLANK_LABEL, zero_infinity=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    best_acc = 0
    early_stop_acc = 0.90

    for epoch in range(50):  # 最多训练50轮
        model.train()
        total_loss = 0
        for images, targets, input_lens, target_lens in train_loader:
            images, targets = images.to(device), targets.to(device)
            logits = model(images)  # [B, T, C]
            log_probs = logits.permute(1, 0, 2)
            loss = criterion(log_probs, targets, input_lens, target_lens)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        acc = evaluate(model, val_loader, device)
        print(f"Epoch {epoch + 1} Loss: {total_loss:.4f}, ValAcc: {acc:.4f}")

        if acc > best_acc:
            best_acc = acc
            torch.save(model.state_dict(), "captcha_crnn_best.pth")
            print("✅ Best model saved.")

        if acc >= early_stop_acc:
            print(f"🎉 Early stopped at epoch {epoch + 1}, val accuracy {acc:.4f}")
            break

    torch.save(model.state_dict(), "captcha_crnn.pth")

    # 导出 ONNX
    dummy = torch.randn(1, 1, 40, 100).to(device)
    torch.onnx.export(model, dummy, "captcha_crnn.onnx",
                      input_names=['input'], output_names=['output'],
                      dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch', 1: 'seq'}},
                      opset_version=11)


def decode_batch(logits_batch):
    preds = logits_batch.argmax(dim=2)  # [B, T]
    results = []
    for pred in preds:
        decoded = []
        prev = -1
        for p in pred.cpu().numpy():
            if p != prev and p != BLANK_LABEL:
                decoded.append(idx2char.get(p, ''))
            prev = p
        results.append(''.join(decoded))
    return results


def evaluate(model, val_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, targets, input_lens, target_lens in val_loader:
            images = images.to(device)
            outputs = model(images)  # [B, T, C]
            preds = decode_batch(outputs)
            labels = targets.cpu().numpy().tolist()
            label_lens = target_lens.cpu().numpy().tolist()

            start = 0
            true_texts = []
            for l in label_lens:
                true_texts.append(''.join([idx2char.get(i, '') for i in labels[start:start + l]]))
                start += l

            for p, t in zip(preds, true_texts):
                if p == t:
                    correct += 1
                total += 1
    return correct / total if total > 0 else 0.0


if __name__ == "__main__":
    train()
